import komand
from .schema import ExecuteExploitInput, ExecuteExploitOutput
# Custom imports below
from metasploit.msfrpc import MsfConsole
import time


class ExecuteExploit(komand.Action):

    def __init__(self):
        super(self.__class__, self).__init__(
                name='execute_exploit',
                description='Run selected Metasploit exploit',
                input=ExecuteExploitInput(),
                output=ExecuteExploitOutput())

    def run(self, params={}):
        # Initial Setup
        client = self.connection.client
        console = MsfConsole(client)
        module = params["module"]
        options = params["options"]
        # Clear console output
        console.read()
        # Set Module
        console.write("use {}".format(module))
        # Iterate through options and set values
        for option, value in options.items():
            console.write("set {} {}".format(option, value))
        # Run module
        console.write("exploit -z")
        # Collect Console Output
        msfReady = False
        console_output = []
        session_info = {}
        try:
            while not msfReady:
                # Added wait for exploit to finish executing
                time.sleep(2)
                msfResult = console.read()
                # Looking for busy key in output
                if (not msfResult['busy']):
                    msfReady = True
                # parsing output
                output = msfResult["data"]
                output_list = output.split("\n")
                output_ids = ["[+]", "[*]", "[!]"]
                for data in output_list:
                    for ids in output_ids:
                        if ids in data:
                            # Strip IDs and add to list
                            console_output.append(data[4:])

                # Collector for Session
                msf_sessions = client.sessions.list
                for msf_session in msf_sessions:
                    target_host = msf_sessions[msf_session]["target_host"]
                    via_exploit = msf_sessions[msf_session]["via_exploit"]
                    if "RHOST" in options:
                        if target_host == options["RHOST"] and via_exploit == module:
                            # Extend session information to include session number
                            msf_sessions[msf_session]["session_id"] = msf_session
                            session_info = msf_sessions[msf_session]
        except Exception as e:
            self.logger.error("An error occurred while executing exploit: ", e)
            raise

        return {"console_output":console_output, "session_information":session_info}

    def test(self):
        client = self.connection.client
        try:
            exploits = client.modules.exploits
        except Exception as e:
            self.logger.error("Error occurred while testing search for exploits: ", e)
            raise
        return {"success": True}
